"""WebSocket连接管理器."""

import asyncio
import json
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Set
from uuid import uuid4

from fastapi import WebSocket, WebSocketDisconnect
from starlette.websockets import WebSocketState

logger = logging.getLogger(__name__)


class WebSocketConnection:
    """WebSocket连接包装器."""
    
    def __init__(self, websocket: WebSocket, client_id: str, user_id: str):
        self.websocket = websocket
        self.client_id = client_id
        self.user_id = user_id
        self.task_id: Optional[str] = None
        self.connected_at = datetime.utcnow()
        self.last_heartbeat = datetime.utcnow()
        self.is_authenticated = False
        
    async def send_json(self, data: dict) -> bool:
        """发送JSON数据."""
        try:
            if self.websocket.client_state == WebSocketState.CONNECTED:
                await self.websocket.send_json(data)
                return True
        except Exception as e:
            logger.error(f"发送消息失败 {self.client_id}: {e}")
        return False
        
    async def send_text(self, text: str) -> bool:
        """发送文本数据."""
        try:
            if self.websocket.client_state == WebSocketState.CONNECTED:
                await self.websocket.send_text(text)
                return True
        except Exception as e:
            logger.error(f"发送文本失败 {self.client_id}: {e}")
        return False


class WebSocketManager:
    """WebSocket连接管理器."""
    
    def __init__(self):
        # 活跃连接映射: client_id -> connection
        self._connections: Dict[str, WebSocketConnection] = {}
        # 用户连接映射: user_id -> Set[client_id]
        self._user_connections: Dict[str, Set[str]] = {}
        # 任务连接映射: task_id -> Set[client_id]
        self._task_connections: Dict[str, Set[str]] = {}
        # 心跳检测任务
        self._heartbeat_task: Optional[asyncio.Task] = None
        # 配置
        self.heartbeat_interval = 30  # 秒
        self.heartbeat_timeout = 60  # 秒
        self.max_connections_per_user = 5
        
    async def start(self):
        """启动管理器."""
        if not self._heartbeat_task:
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
            logger.info("WebSocket管理器已启动")
            
    async def stop(self):
        """停止管理器."""
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
            except asyncio.CancelledError:
                pass
            self._heartbeat_task = None
            
        # 断开所有连接
        for client_id in list(self._connections.keys()):
            await self.disconnect(client_id)
        logger.info("WebSocket管理器已停止")
        
    async def connect(self, websocket: WebSocket, user_id: str) -> str:
        """建立WebSocket连接.
        
        Args:
            websocket: WebSocket对象
            user_id: 用户ID
            
        Returns:
            客户端ID
            
        Raises:
            ValueError: 如果用户连接数超限
        """
        # 检查用户连接数
        user_connections = self._user_connections.get(user_id, set())
        if len(user_connections) >= self.max_connections_per_user:
            raise ValueError(f"用户 {user_id} 连接数已达上限")
            
        # 接受连接
        await websocket.accept()
        
        # 生成客户端ID
        client_id = str(uuid4())
        
        # 创建连接对象
        connection = WebSocketConnection(websocket, client_id, user_id)
        connection.is_authenticated = True  # 假设已通过JWT认证
        
        # 保存连接
        self._connections[client_id] = connection
        
        # 更新用户连接映射
        if user_id not in self._user_connections:
            self._user_connections[user_id] = set()
        self._user_connections[user_id].add(client_id)
        
        # 发送连接成功消息
        await connection.send_json({
            "type": "connection",
            "client_id": client_id,
            "timestamp": datetime.utcnow().isoformat()
        })
        
        logger.info(f"WebSocket连接建立: client={client_id}, user={user_id}")
        return client_id
        
    async def disconnect(self, client_id: str):
        """断开WebSocket连接.
        
        Args:
            client_id: 客户端ID
        """
        connection = self._connections.get(client_id)
        if not connection:
            return
            
        # 从任务连接中移除
        if connection.task_id:
            task_connections = self._task_connections.get(connection.task_id, set())
            task_connections.discard(client_id)
            if not task_connections:
                del self._task_connections[connection.task_id]
                
        # 从用户连接中移除
        user_connections = self._user_connections.get(connection.user_id, set())
        user_connections.discard(client_id)
        if not user_connections:
            del self._user_connections[connection.user_id]
            
        # 关闭WebSocket
        try:
            if connection.websocket.client_state == WebSocketState.CONNECTED:
                await connection.websocket.close()
        except Exception as e:
            logger.error(f"关闭WebSocket失败 {client_id}: {e}")
            
        # 移除连接
        del self._connections[client_id]
        logger.info(f"WebSocket连接断开: client={client_id}")
        
    async def bind_task(self, client_id: str, task_id: str):
        """绑定客户端到任务.
        
        Args:
            client_id: 客户端ID
            task_id: 任务ID
        """
        connection = self._connections.get(client_id)
        if not connection:
            raise ValueError(f"客户端 {client_id} 不存在")
            
        # 解绑旧任务
        if connection.task_id:
            old_task_connections = self._task_connections.get(connection.task_id, set())
            old_task_connections.discard(client_id)
            
        # 绑定新任务
        connection.task_id = task_id
        if task_id not in self._task_connections:
            self._task_connections[task_id] = set()
        self._task_connections[task_id].add(client_id)
        
        logger.info(f"客户端 {client_id} 绑定到任务 {task_id}")
        
    async def send_to_client(self, client_id: str, message: dict) -> bool:
        """发送消息到指定客户端.
        
        Args:
            client_id: 客户端ID
            message: 消息内容
            
        Returns:
            是否发送成功
        """
        connection = self._connections.get(client_id)
        if connection:
            message["timestamp"] = datetime.utcnow().isoformat()
            return await connection.send_json(message)
        return False
        
    async def send_to_task(self, task_id: str, message: dict) -> int:
        """发送消息到任务的所有客户端.
        
        Args:
            task_id: 任务ID
            message: 消息内容
            
        Returns:
            成功发送的客户端数
        """
        client_ids = self._task_connections.get(task_id, set()).copy()
        message["timestamp"] = datetime.utcnow().isoformat()
        
        success_count = 0
        for client_id in client_ids:
            if await self.send_to_client(client_id, message):
                success_count += 1
                
        return success_count
        
    async def send_to_user(self, user_id: str, message: dict) -> int:
        """发送消息到用户的所有客户端.
        
        Args:
            user_id: 用户ID
            message: 消息内容
            
        Returns:
            成功发送的客户端数
        """
        client_ids = self._user_connections.get(user_id, set()).copy()
        message["timestamp"] = datetime.utcnow().isoformat()
        
        success_count = 0
        for client_id in client_ids:
            if await self.send_to_client(client_id, message):
                success_count += 1
                
        return success_count
        
    async def broadcast(self, message: dict) -> int:
        """广播消息到所有客户端.
        
        Args:
            message: 消息内容
            
        Returns:
            成功发送的客户端数
        """
        message["timestamp"] = datetime.utcnow().isoformat()
        
        success_count = 0
        for client_id in list(self._connections.keys()):
            if await self.send_to_client(client_id, message):
                success_count += 1
                
        return success_count
        
    async def handle_message(self, client_id: str, message: dict):
        """处理客户端消息.
        
        Args:
            client_id: 客户端ID
            message: 消息内容
        """
        connection = self._connections.get(client_id)
        if not connection:
            return
            
        message_type = message.get("type")
        
        if message_type == "heartbeat":
            # 更新心跳时间
            connection.last_heartbeat = datetime.utcnow()
            await connection.send_json({
                "type": "heartbeat_ack",
                "timestamp": datetime.utcnow().isoformat()
            })
            
        elif message_type == "bind_task":
            # 绑定任务
            task_id = message.get("task_id")
            if task_id:
                await self.bind_task(client_id, task_id)
                await connection.send_json({
                    "type": "task_bound",
                    "task_id": task_id,
                    "timestamp": datetime.utcnow().isoformat()
                })
                
        else:
            logger.warning(f"未知消息类型: {message_type}")
            
    async def _heartbeat_loop(self):
        """心跳检测循环."""
        while True:
            try:
                await asyncio.sleep(self.heartbeat_interval)
                await self._check_heartbeats()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"心跳检测异常: {e}")
                
    async def _check_heartbeats(self):
        """检查心跳超时."""
        now = datetime.utcnow()
        timeout_threshold = now - timedelta(seconds=self.heartbeat_timeout)
        
        # 发送心跳请求
        await self.broadcast({"type": "heartbeat_request"})
        
        # 检查超时连接
        timeout_clients = []
        for client_id, connection in self._connections.items():
            if connection.last_heartbeat < timeout_threshold:
                timeout_clients.append(client_id)
                
        # 断开超时连接
        for client_id in timeout_clients:
            logger.warning(f"客户端 {client_id} 心跳超时，断开连接")
            await self.disconnect(client_id)
            
    def get_connection_info(self) -> dict:
        """获取连接统计信息."""
        return {
            "total_connections": len(self._connections),
            "total_users": len(self._user_connections),
            "total_tasks": len(self._task_connections),
            "connections_by_user": {
                user_id: len(clients) 
                for user_id, clients in self._user_connections.items()
            }
        }


# 全局WebSocket管理器实例
websocket_manager = WebSocketManager()